# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import json
import logging
import math
import os
import re
import statistics
from collections import OrderedDict

import pprint
import tabulate
import pandas

from coin.base.param_util import to_list


def load_sim_result(result_dirs, name_filter=None):
  result_dirs = to_list(result_dirs)
  name_filter = name_filter or (lambda _: True)
  if isinstance(name_filter, str):
    pattern = name_filter
    name_filter = lambda name: re.search(pattern, name)

  results = {}
  for res_dir in result_dirs:
    res_dir = os.path.normpath(res_dir)
    for filename in sorted(os.listdir(res_dir)):
      filepath = os.path.join(res_dir, filename)
      if (not filename.startswith('result-') or not filename.endswith('.json')
          or not os.path.isfile(filepath)):
        continue

      with open(filepath, 'r') as f:
        result_subset = json.load(f)

      for res in result_subset:
        if name_filter(res['name']):
          results[res['name']] = res  # Overwrite duplicate with later one.

  return sorted(results.values(), key=lambda x: x['name'])


def aggregate_sim_result(results, aggregate_key_func=None, show_stat_only=True):
  results = sorted(results, key=lambda r: r['name'])

  if not show_stat_only:
    pprint.pprint(results)
    print('---')

    try:
      headers = ['name', 'num_trade', 'qty_sum', 'pq_sum', 'pnl', 'return_bp', 'fee', 'fee_bp']
      rows = [[r[h] for h in headers] for r in results]
      print(tabulate.tabulate(rows, headers=headers))
    except KeyError:
      pass
    except Exception:
      logging.exception('Exception occured')
    print('---')

  print_sim_result_stat(results, aggregate_key_func)


def print_sim_result_stat(results, key_func=None):
  # Assumes results is
  # (1) list of dicts
  # (2) each dict has 'name' key
  # (3) last 8 chracter of 'name' key should be date in YYYYMMDD format.

  if key_func is None:
    key_func = lambda r: (r['name'][:-9], r['name'][-8:])

  times = set()
  agg = {}
  for r in results:
    param_str, time_str = key_func(r)
    times.add(time_str)
    if param_str not in agg:
      agg[param_str] = {}
    if time_str not in agg[param_str]:
      agg[param_str][time_str] = []
    agg[param_str][time_str].append(r)

  params = sorted(agg.keys())
  df_dict = OrderedDict(time=[])
  for p in params:
    df_dict[p] = []

  pnl = {p: [] for p in params}
  qty = {p: [] for p in params}
  cnt = {p: [] for p in params}
  fee = {p: [] for p in params}
  sub_qty = {p: [] for p in params}
  pq_dict = {p: [] for p in params}
  for date_str in sorted(times):
    df_dict['time'].append(date_str)

    for param in params:
      p = sum([r['pnl'] for r in agg[param][date_str]])
      q = sum([r['qty_sum'] for r in agg[param][date_str]])
      ntrade = sum([r['num_trade'] for r in agg[param][date_str]])
      feesum = sum([r['fee'] for r in agg[param][date_str]])
      subqsum = sum([r['sub_qty_sum'] for r in agg[param][date_str]])

      pnl[param].append(p)
      qty[param].append(q)
      cnt[param].append(ntrade)
      df_dict[param].append(p)
      fee[param].append(feesum)
      sub_qty[param].append(subqsum)
      try:
        pq = sum([r['pq_sum'] for r in agg[param][date_str]])
        pq_dict[param].append(pq)
      except Exception:
        pass

  df_dict['time'] += [
      '---', 'mean', 'stdev', 'sharpe', 'cnt_mean', 'qty_mean', 'ret_bp', 'fee_bp', 'f/q'
  ]
  for param in params:
    pnl_sum = sum(pnl[param])
    if pnl[param] and qty[param]:
      mean = statistics.mean(pnl[param])
      if len(pnl[param]) > 1:
        stdev = statistics.stdev(pnl[param])
        sharpe = mean / stdev if stdev != 0 else None
      else:
        stdev = None
        sharpe = None
      qty_mean = statistics.mean(qty[param])
      cnt_mean = statistics.mean(cnt[param])
    else:
      mean = None
      stdev = None
      sharpe = None
      qty_mean = None
      cnt_mean = None

    if len(pq_dict[param]) == 0:
      ret_bp = None
      fee_bp = None
      fr_qty = None
    elif sum(pq_dict[param]) != 0:
      ret_bp = sum(pnl[param]) / sum(pq_dict[param]) * 1e4
      fee_bp = sum(fee[param]) / sum(pq_dict[param]) * 1e4
      fr_qty = sum(qty[param]) / sum(sub_qty[param])
    else:
      ret_bp = 0.
      fee_bp = 0.
      fr_qty = 0.
    df_dict[param] += [None, mean, stdev, sharpe, cnt_mean, qty_mean, ret_bp, fee_bp, fr_qty]

  df = pandas.DataFrame(df_dict)
  df = df.set_index('time')
  pandas.options.display.float_format = lambda x: '' if math.isnan(x) else '{:,.2f}'.format(x)
  with pandas.option_context('display.max_rows',
                             None,
                             'display.max_columns',
                             None,
                             'display.width',
                             250):
    print(df)
